# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
WithEnergyCell
"""

from typing import Union, List, Tuple
import torch
from torch import Tensor
from torch.nn import Parameter, Module, ModuleList


from ...function import Units, get_arguments
from ...partition import NeighbourList
from ...system import Molecule
from ...potential import PotentialCell
from ...potential.bias import Bias
from ...sampling.wrapper import EnergyWrapper
from ...function import GLOBAL_DEVICE_SPONGE


class WithEnergyCell(Module):

    def __init__(self,
                 system: Molecule,
                 potential: PotentialCell,
                 bias: Union[Bias, List[Bias]] = None,
                 cutoff: float = None,
                 neighbour_list: NeighbourList = None,
                 wrapper: EnergyWrapper = None,
                 **kwargs
                 ):

        super().__init__()
        self._kwargs = get_arguments(locals(), kwargs)

        self.device = GLOBAL_DEVICE_SPONGE()
        self.system = system
        self.potential_function = potential

        self.units = Units(self.system.length_unit, self.potential_function.energy_unit)
        self.system.units.set_energy_unit(self.energy_unit)

        self.bias_function: List[Bias] = None
        self._num_biases = 0
        self._bias_names = []
        if bias is not None:
            if isinstance(bias, list):
                self._num_biases = len(bias)
                self.bias_function = ModuleList(bias)
            elif isinstance(bias, Module):
                self._num_biases = 1
                self.bias_function = ModuleList([bias])
            else:
                raise TypeError(f'The "bias" must be Cell or list but got: {type(bias)}')

            for i in range(self._num_biases):
                self._bias_names.append(self.bias_function[i].name)

        self.num_walker = self.system.num_walker
        self.num_atoms = self.system.num_atoms

        self.energy_wrapper = wrapper
        if wrapper is None:
            self.energy_wrapper = EnergyWrapper(length_unit=self.length_unit, energy_unit=self.energy_unit)

        self.wrapper_pace = self.energy_wrapper.update_pace

        self.exclude_index = self.potential_function.exclude_index
        self.neighbour_list = neighbour_list
        if neighbour_list is None:
            if cutoff is None and self.potential_function.cutoff is not None:
                cutoff = self.potential_function.cutoff
            self.neighbour_list = NeighbourList(system, cutoff, exclude_index=self.exclude_index)
        else:
            self.neighbour_list.set_exclude_index(self.exclude_index)

        self.neighbour_index = self.neighbour_list.neighbours
        self.neighbour_mask = self.neighbour_list.neighbour_mask
        self.num_neighbours = self.neighbour_list.num_neighbours

        if self.neighbour_list.cutoff is not None:
            if self.potential_function.cutoff is None:
                self.potential_function.set_cutoff(self.neighbour_list.cutoff, self.length_unit)
            else:
                pot_cutoff = self.units.length(self.potential_function.cutoff, self.potential_function.length_unit)
                nl_cutoff = self.neighbour_list.cutoff
                if self.potential_function.cutoff > self.neighbour_list.cutoff:
                    raise ValueError(f'The cutoff of the potential function ({pot_cutoff} {self.length_unit}) '
                                     f'cannot be greater than '
                                     f'the cutoff of the neighbour list ({nl_cutoff} {self.length_unit}).')

        self.coordinate = self.system.coordinate.to(self.device)
        self.pbc_box = self.system.pbc_box.to(self.device) if self.system.pbc_box is not None else None
        self.atom_mass = self.system.atom_mass

        self.potential_function.set_pbc(self.pbc_box is not None)

        for p in self.potential_function.parameters():
            p.requires_grad = False

        self.potential_function_units = self.potential_function.units

        self.length_unit_scale = torch.tensor(self.units.convert_length_to(
            self.potential_function.length_unit), dtype=torch.float32)

        self._energies = Parameter(torch.zeros((self.num_walker, self.num_energies),
                                              dtype=torch.float32), requires_grad=False)

        bias = torch.zeros((self.num_walker, 1), dtype=torch.float32)
        if self.bias_function is None:
            self._biases = None
            self._bias = bias
        else:
            self._biases = Parameter(torch.zeros((self.num_walker, self._num_biases),
                                                dtype=torch.float32), requires_grad=False)
            self._bias = Parameter(bias, requires_grad=False)

    @property
    def cutoff(self) -> Tensor:

        return self.neighbour_list.cutoff

    @property
    def neighbour_list_pace(self) -> int:

        return self.neighbour_list.pace

    @property
    def length_unit(self) -> str:

        return self.units.length_unit

    @property
    def energy_unit(self) -> str:

        return self.units.energy_unit

    @property
    def num_energies(self) -> int:

        return self.potential_function.num_energies

    @property
    def num_biases(self) -> int:

        return self._num_biases

    @property
    def energy_names(self) -> list:

        return self.potential_function.energy_names

    @property
    def bias_names(self) -> list:

        return self._bias_names

    @property
    def energies(self) -> Tensor:

        return self._energies.detach()

    @property
    def biases(self) -> Tensor:

        if self.bias_function is None:
            return None
        return self._biases.detach()

    @property
    def bias(self) -> Tensor:
   
        return self._bias.detach()

    def bias_pace(self, index: int = 0) -> int:
    
        return self.bias_function[index].update_pace

    def set_pbc_grad(self, grad_box: bool):
  
        self.system.set_pbc_grad(grad_box)
        return self

    def update_neighbour_list(self) -> Tuple[Tensor, Tensor]:
 
        return self.neighbour_list.update(self.coordinate, self.pbc_box)

    def update_bias(self, step: int):
    
        if self.bias_function is not None:
            for i in range(self._num_biases):
                if self.bias_pace(i) > 0 and step % self.bias_pace(i) == 0:
                    self.bias_function[i].update(self.coordinate, self.pbc_box)
        return self

    def update_wrapper(self, step: int):
     
        if self.wrapper_pace > 0 and step % self.wrapper_pace == 0:
            self.energy_wrapper.update()
        return self

    def get_neighbour_list(self) -> Tuple[Tensor, Tensor]:
        
        return self.neighbour_list.get_neighbour_list()

    def calc_energies(self) -> Tensor:
   
        neigh_idx, neigh_vec, neigh_dis, neigh_mask = self.neighbour_list(self.coordinate, self.pbc_box)

        coordinate = self.coordinate * self.length_unit_scale
        pbc_box = self.pbc_box
        if pbc_box is not None:
            pbc_box = pbc_box * self.length_unit_scale
        else:
            pbc_box = None
        neigh_vec = neigh_vec * self.length_unit_scale
        neigh_dis = neigh_dis * self.length_unit_scale

        energies = self.potential_function(
            coordinate=coordinate,
            neighbour_index=neigh_idx,
            neighbour_mask=neigh_mask,
            neighbour_vector=neigh_vec,
            neighbour_distance=neigh_dis,
            pbc_box=pbc_box
        )

        return energies

    def calc_biases(self) -> Tensor:
        """
        Calculate the bias potential terms.

        Return:
            Tensor, Tensor of shape `(B, V)`. Data type is float. Bias potential terms.

        Symbols:
            B:  Batchsize, i.e. number of walkers of the simulation.
            V:  Number of bias potential terms.
        """
        if self.bias_function is None:
            return None

        neigh_idx, neigh_vec, neigh_dis, neigh_mask = self.neighbour_list(self.coordinate, self.pbc_box)

        coordinate = self.coordinate * self.length_unit_scale
        if self.pbc_box is not None:
            pbc_box = self.pbc_box * self.length_unit_scale
        else:
            pbc_box = None
        neigh_vec = neigh_vec * self.length_unit_scale
        neigh_dis = neigh_dis * self.length_unit_scale

        biases = ()
        for i in range(self._num_biases):
            bias_ = self.bias_function[i](
                coordinate=coordinate,
                neighbour_index=neigh_idx,
                neighbour_mask=neigh_mask,
                neighbour_vector=neigh_vec,
                neighbour_distance=neigh_dis,
                pbc_box=pbc_box)
            # bias_ = self.bias_function[i].colvar(
            #     coordinate=coordinate,
            #     pbc_box=pbc_box
            # )
            biases += (bias_,)

        return torch.cat(biases, dim=-1)

    def print_info(self): 
        if self.potential_function is not None:
            print("[TorchSPONGE] Energy terms include:\n[TorchSPONGE] ",end='')
            for i, n in enumerate(self.potential_function._energy_names):
                print(f"{i}. {n}",end='    ')
            print(f"\n{'--'*40}")

        if self.bias_function is not None:
            for i in range(self._num_biases):
                self.bias_function[i].print_info()
    
    def forward(self, *inputs) -> Tensor:
        """calculate the total potential energy (potential energy and bias potential) of the simulation system.

        Return:
            energy (Tensor):    Tensor of shape `(B, 1)`. Data type is float.
                                Total potential energy.

        Symbols:
            B:  Batchsize, i.e. number of walkers of the simulation.

        """
        #pylint: disable=unused-argument
        coordinate, pbc_box = self.system()
        coordinate = coordinate.to(self.device)
        pbc_box = pbc_box.to(self.device) if pbc_box is not None else None

        neigh_idx, neigh_vec, neigh_dis, neigh_mask = self.neighbour_list(coordinate, pbc_box)

        coordinate = coordinate * self.length_unit_scale
        if pbc_box is not None:
            pbc_box = pbc_box * self.length_unit_scale

        if neigh_idx is not None:
            neigh_vec = neigh_vec * self.length_unit_scale
            neigh_dis = neigh_dis * self.length_unit_scale

        energies = self.potential_function(
            coordinate=coordinate,
            neighbour_index=neigh_idx,
            neighbour_mask=neigh_mask,
            neighbour_vector=neigh_vec,
            neighbour_distance=neigh_dis,
            pbc_box=pbc_box
        )

        self._energies.data.copy_(energies)

        biases = None
        if self.bias_function is not None:
            biases = ()
            for i in range(self._num_biases):
                bias_ = self.bias_function[i](
                    coordinate=coordinate,
                    neighbour_index=neigh_idx,
                    neighbour_mask=neigh_mask,
                    neighbour_vector=neigh_vec,
                    neighbour_distance=neigh_dis,
                    pbc_box=pbc_box
                )
                biases += (bias_,)

            biases = torch.cat(biases, dim=-1)
            self._biases.data.copy_(biases)

        energy, bias = self.energy_wrapper(energies, biases)

        if self.bias_function is not None:
            self._bias.data.copy_(bias)

        return energy
